package com.ocom.amqp.core;

import com.ocom.amqp.builder.AbstractAssemblyBuilder;
import lombok.Data;
import org.springframework.amqp.AmqpException;
import org.springframework.amqp.core.Message;
import org.springframework.amqp.core.MessagePostProcessor;
import org.springframework.amqp.core.MessageProperties;
import org.springframework.amqp.rabbit.core.RabbitAdmin;

import java.util.Properties;

@Data
public class CurrentLimitingQueue extends BasicExpandAssembly {

    private String businessName;

    private String routingKey;

    //限流数量
    private Integer currentLimiting;

    //限流时间单位
    private Long timeUnit;

    public CurrentLimitingQueue(String businessName, String routingKey, Integer currentLimiting, Long timeUnit) {
        this.businessName = businessName;
        this.routingKey = routingKey;
        this.currentLimiting = currentLimiting;
        this.timeUnit = timeUnit;
    }

    /**
     * 限流器 在{@link com.ocom.amqp.builder.CurrentLimitingQueueBuilder} 中
     * 进行注入
     */
    private CurrentLimiter currentLimiter;


    public String getBufferQueueName(){
        return AbstractAssemblyBuilder.DELAY_PREFIX.concat(businessName).concat(AbstractAssemblyBuilder.QUEUE_SUFFIX);
    }

    @Override
    public String getSendExchangeName() {
        return AbstractAssemblyBuilder.DELAY_PREFIX.concat(businessName).concat(AbstractAssemblyBuilder
                .EXCHANGE_SUFFIX);
    }

    @Override
    public String getListenQueueName() {
        return AbstractAssemblyBuilder.DEAD_PREFIX.concat(businessName).concat(AbstractAssemblyBuilder.QUEUE_SUFFIX);
    }

    /**
     * 限流器
     */
    public static class CurrentLimiter implements MessagePostProcessor {

        private CurrentLimitingQueue assembly;

        private RabbitAdmin rabbitAdmin;

        private Long preTime = 0l;

        public CurrentLimiter(CurrentLimitingQueue currentLimitingQueue, RabbitAdmin rabbitAdmin) {
            this.assembly = currentLimitingQueue;
            this.rabbitAdmin = rabbitAdmin;
            this.preTime = (long) Math.ceil(currentLimitingQueue.getTimeUnit() / currentLimitingQueue
                    .getCurrentLimiting());
        }

        @Override
        public Message postProcessMessage(Message message) throws AmqpException {
            MessageProperties messageProperties = message.getMessageProperties();
            String receivedExchange = messageProperties.getReceivedExchange();
            String receivedRoutingKey = messageProperties.getReceivedRoutingKey();
            if (receivedExchange.equals(assembly.getSendExchangeName()) && receivedRoutingKey
                    .equals(assembly.getRoutingKey())) {
                //如果是发送到此限流模型中 进行限流
                Long expiration = this.calculateExpiration(assembly.getCurrentLimiting(), assembly.getTimeUnit(),
                        assembly
                                .getBufferQueueName());
                message.getMessageProperties().setExpiration(String.valueOf(expiration));
            }
            return message;
        }

        /**
         * 计算此消息超时时间应设为多少
         *
         * @param limitingCount
         * @param timeUnit
         * @return
         */
        private Long calculateExpiration(Integer limitingCount, Long timeUnit, String queueName) {
            Properties queueProperties = this.rabbitAdmin.getQueueProperties(queueName);
            Integer messageCount = (Integer) queueProperties.get(RabbitAdmin.QUEUE_MESSAGE_COUNT);
            return this.preTime * (messageCount + 1);
        }
    }
}
